import os
import cv2
from tqdm import tqdm
from glob import glob
from osgeo import gdal, osr


def EPSG4326_to_EPSG3857(latitude, longitude):
    """将WGS 84坐标系下的点转换到Web Mercator坐标系"""
    # 创建WGS 84坐标系
    wgs84_crs = osr.SpatialReference()
    wgs84_crs.ImportFromEPSG(4326)

    # 创建Web Mercator坐标系
    web_mercator_crs = osr.SpatialReference()
    web_mercator_crs.ImportFromEPSG(3857)

    # 创建坐标变换对象
    coord_transform = osr.CoordinateTransformation(wgs84_crs, web_mercator_crs)

    # 将经纬度坐标转换为Web Mercator坐标
    # web_mercator_x, web_mercator_y, _ = coord_transform.TransformPoint(longitude, latitude)
    web_mercator_x, web_mercator_y, _ = coord_transform.TransformPoint(latitude,longitude)

    return web_mercator_x, web_mercator_y


def EPSG3857_to_EPSG4326(latitude, longitude):
    # 创建WGS 84坐标系
    wgs84_crs = osr.SpatialReference()
    wgs84_crs.ImportFromEPSG(4326)

    # 创建Web Mercator坐标系
    web_mercator_crs = osr.SpatialReference()
    web_mercator_crs.ImportFromEPSG(3857)

    # 创建坐标变换对象
    coord_transform = osr.CoordinateTransformation(web_mercator_crs, wgs84_crs)

    wgs84_y, wgs84_x, _ = coord_transform.TransformPoint(longitude, latitude)
    
    return wgs84_x, wgs84_y


def geo2pixel(geo_transform,lon,lat):
    # 将经纬度转换为图像的像素坐标
    col = int((lon - geo_transform[0]) / geo_transform[1])
    row = int((lat - geo_transform[3]) / geo_transform[5])
    return col, row


def pixel_to_geo_coords(geo_transform, px, py):
    """从像素坐标转换到地理坐标"""
    ulx, xres, xskew, uly, yskew, yres = geo_transform
    lon = ulx + (px + 1) * xres + py * xskew
    lat = uly + (px + 1) * yskew + py * yres
    return lon, lat


def cut_image(src_ds, x_min, y_min, x_max, y_max):


    # 获取图像的地理变换矩阵
    geo_transform = src_ds.GetGeoTransform()
    if geo_transform is None:
        print("Failed to get GeoTransform from input file.")
        return


    # 裁剪区域宽度和高度
    width = x_max - x_min
    height = y_max - y_min
    lon, lat = pixel_to_geo_coords(geo_transform,x_min, y_min)
    # 创建一个新的GeoTransform矩阵
    new_geo_transform = list(geo_transform)
    new_geo_transform[0] = lon
    new_geo_transform[3] = lat

    # 创建一个内存中的临时数据集用于裁剪
    mem_drv = gdal.GetDriverByName('MEM')
    dst_ds = mem_drv.Create('', width, height, src_ds.RasterCount, src_ds.GetRasterBand(1).DataType)
    if dst_ds is None:
        print('Failed to create memory raster.')
        return

    # 设置新的地理变换矩阵
    dst_ds.SetGeoTransform(new_geo_transform)

    # 复制投影信息
    dst_ds.SetProjection(src_ds.GetProjection())

    # 执行裁剪
    # 注意: 这里使用src_ds.ReadAsArray方法直接从源数据集中读取指定范围的数据
    data = src_ds.ReadAsArray(x_min, y_min, width, height)

    for band_num in range(1, src_ds.RasterCount + 1):
        band = dst_ds.GetRasterBand(band_num)
        band.WriteArray(data[band_num - 1])
    del src_ds

    return dst_ds


def save_tif(output_file,dst_ds):
    if output_file:
    # 将内存中的图像保存到磁盘
        out_driver = gdal.GetDriverByName('GTiff')
        out_data = out_driver.CreateCopy(output_file, dst_ds, 0)
        return out_data
    else:
        print('output_file is False')


def read_img(path):
    image = gdal.Open(path)

    lon, lat = image.GetGeoTransform()[0], image.GetGeoTransform()[3]
    width, height = image.RasterXSize, image.RasterYSize
    trans = image.GetGeoTransform()
    bottom_right_x = trans[0] + trans[1] * width + trans[2] * height
    bottom_right_y = trans[3] + trans[4] * width + trans[5] * height
    return image, lon, lat, bottom_right_x, bottom_right_y


if __name__ == '__main__':
    '''
    影像预处理，根据Google和Bing地图的地理坐标，使两张图像强制对齐，误差小于两像素
    处理后Bing图像用作制作基准图像，Google图像用作制作查询图像
    '''
    root_path = r'D:/Temp'

    google_map_path = f'{root_path}/google'
    bing_map_path = f'{root_path}/bing'

    image_list = glob(f'{google_map_path}/*.tif')
    loop = tqdm((image_list), total = len(image_list))
    for img_path in loop:
        img_name = os.path.basename(img_path)
        google_image, google_lon, google_lat, google_bottom_right_lon, google_bottom_right_lat = read_img(f'{google_map_path}/{img_name}')
        bing_image, bing_lon, bing_lat, bing_bottom_right_lon, bing_bottom_right_lat = read_img(f'{bing_map_path}/{img_name}')

        google_trans = google_image.GetGeoTransform()
        bing_trans = bing_image.GetGeoTransform()

        bing_lon, bing_lat = EPSG4326_to_EPSG3857(bing_lat, bing_lon)
        bing_bottom_right_lon, bing_bottom_right_lat = EPSG4326_to_EPSG3857(bing_bottom_right_lat, bing_bottom_right_lon)

        tp_lon, tp_lat = max(google_lon, bing_lon), min(google_lat, bing_lat)
        br_lon, br_lat = min(google_bottom_right_lon, bing_bottom_right_lon), max(google_bottom_right_lat, bing_bottom_right_lat)

        google_tp_col, google_tp_row = geo2pixel(google_trans, tp_lon, tp_lat)
        google_br_col, google_br_row = geo2pixel(google_trans, br_lon, br_lat)

        bing_tp_lon, bing_tp_lat = EPSG3857_to_EPSG4326(tp_lat, tp_lon)
        bing_br_lon, bing_br_lat = EPSG3857_to_EPSG4326(br_lat, br_lon)

        bing_tp_col, bing_tp_row = geo2pixel(bing_trans, bing_tp_lon, bing_tp_lat)
        bing_br_col, bing_br_row = geo2pixel(bing_trans, bing_br_lon, bing_br_lat)

        bing_img = cut_image(bing_image, bing_tp_col, bing_tp_row, bing_br_col, bing_br_row)
        google_img = cut_image(google_image, google_tp_col, google_tp_row, google_br_col, google_br_row)

        bing_img = bing_img.ReadAsArray()
        bing_img = bing_img.transpose(1, 2, 0)
        google_img = google_img.ReadAsArray()
        google_img = google_img.transpose(1, 2, 0)
        google_img = cv2.resize(google_img, (bing_img.shape[1], bing_img.shape[0]), interpolation=cv2.INTER_CUBIC)
        bing_img = cv2.cvtColor(bing_img, cv2.COLOR_RGB2BGR)
        google_img = cv2.cvtColor(google_img, cv2.COLOR_RGB2BGR)

        os.makedirs(f'{root_path}/bing_process', exist_ok=True)
        os.makedirs(f'{root_path}/google_process', exist_ok=True)

        # save_tif(f'{root_path}/bing_process/{img_name}', bing_img)
        # save_tif(f'{root_path}/google_process/{img_name}', google_img)
        cv2.imwrite(f'{root_path}/bing_process/{img_name}', bing_img)
        cv2.imwrite(f'{root_path}/google_process/{img_name}', google_img)

print('done')